#!/bin/bash

# usage: ./build.sh 22.04-cuda12.1.1-cudnn8-devel

# echo "argv: $1"
UBUNTU_VERSION=`echo $1 | awk -F '-cuda' '{print $1}'`
CUDA_VERSION=`echo $1 | awk -F '-cuda' '{print $2}'`
echo "ubuntu version:${UBUNTU_VERSION},cuda version:${CUDA_VERSION}"

# pull base image (ubuntu/cuda)
if [[("${CUDA_VERSION}" == "")]];then
    BASE_IMAGE=ubuntu:${UBUNTU_VERSION}
else
    BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-ubuntu${UBUNTU_VERSION}
fi
echo ${BASE_IMAGE}

docker pull ${BASE_IMAGE}
if [[ $? != 0 ]]; then
    echo "Failed to pull docker image '${BASE_IMAGE}'"
    exit -2
fi

# build docker-desktop image
if [[("${CUDA_VERSION}" == "")]];then
    DOCKER_TAG=${UBUNTU_VERSION}
else
    DOCKER_TAG=${UBUNTU_VERSION}-cuda${CUDA_VERSION}
fi

docker build ./ --build-arg BASE_IMAGE=${BASE_IMAGE} \
                --tag docker-desktop:${DOCKER_TAG}
if [[ $? != 0 ]]; then
    echo "Failed to build docker image 'docker-desktop:${DOCKER_TAG}'"
    exit -3
fi

exit 0